// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef __Response_hpp__
#define __Response_hpp__

#include <string>

#include "Teuchos_RCP.hpp"
#include "Teuchos_DefaultComm.hpp"

#include "Thyra_VectorSpaceBase.hpp"
#include "Thyra_VectorBase.hpp"
#include "Thyra_MultiVectorBase.hpp"
#include "Thyra_DefaultSpmdVectorSpace.hpp"
#include "Thyra_LinearOpBase.hpp"

#include "ResponseBase.hpp"

namespace SiYuan {

template <typename EvalT>
class Response : public ResponseBase {
public:
  Response(const std::string & responseName)
     : ResponseBase(responseName) 
  {
      tComm_ = Teuchos::DefaultComm<Thyra::Ordinal>::getComm();
  }

  //! What is the number of values you need locally
   virtual std::size_t localSizeRequired() const = 0;

   //! Is the vector distributed (or replicated)
   virtual bool vectorIsDistributed() const = 0;

  //! Get the vector space for this response, vector space is constructed lazily.
   Teuchos::RCP<const Thyra::VectorSpaceBase<double> > getVectorSpace() const;

   /** Set the vector (to be filled) for this response. This must be
     * constructed from the vector space returned by <code>getVectorSpace</code>.
     */
   void setVector(const Teuchos::RCP<Thyra::VectorBase<double> > & destVec)
   { tVector_ = destVec; }

   //! set the vector space for this response
   void setVectorSpace(Teuchos::RCP<const Thyra::VectorSpaceBase<double> > vs)
   { vSpace_ = vs; }

   //! Access the response vector
   Teuchos::RCP<Thyra::VectorBase<double> > getVector() const
   { return tVector_; }

   /** Build the response object used by this factory. This object
     * assumes the role of the scatter target and will be accessible
     * by all the evaluators in the field managers. This is the sideset
     * version of the buildResponseObject function.
     *
     * \param[in] responseName Name of response to be built. This
     *                         name will be used for looking up
     *                         the response in the <code>GlobalEvaluationDataContainer</code>
     *                         object.
     * \param[in] wkstdescs A vector of descriptors for the elements this response is over.
     */
   virtual Teuchos::RCP<ResponseBase> buildResponseObject(const std::string & responseName,
                        const std::vector<panzer::WorksetDescriptor> & wkstdescs) const = 0; 

   /** Build and register evaluators for a response on a particular physics
     * block. Note that it is assumed that a field has been marked required
     * during this method call.
     *
     * \param[in] responseName The name of the response to be constructed
     *                         by these evaluators.
     * \param[in,out] fm Field manager to be fuild with the evaluators.
     * \param[in] physicsBlock What physics block is being used for constructing
     *                         the evaluators
     * \param[in] user_data The user data parameter list, this stores things
     *                      that the user may find useful.
     */
   virtual void buildAndRegisterEvaluators(const std::string & responseName,
                                           PHX::FieldManager<panzer::Traits> & fm,
                                           const panzer::PhysicsBlock & physicsBlock,
                                           const Teuchos::ParameterList & user_data) const = 0;

   /** Is this evaluation type supported by the factory. This is used to determine cases
     * where a response may support a particular evaluation type, however at runtime the user
     * decides not to enable the (say) Jacobian evaluation of this response.
     *
     * Note that use of this mechanism is complementary to having the builder return 
     * <code>Teuchos::null</code> for a particular evaluation type.
     */
   virtual bool typeSupported() const = 0;

protected:
  //! Get the teuchos comm object
  Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > getComm() const { return tComm_; }

  //! Access the thyra vector
  Thyra::ArrayRCP<double> getThyraVector() const;

  //! Access the thyra MultiVector
  Teuchos::RCP<Thyra::MultiVectorBase<double> > getThyraMultiVector() const
  { return tVector_;}

private:
   // hide these methods
   Response();
   Response(const Response<EvalT> &);

   mutable Teuchos::RCP<const Thyra::VectorSpaceBase<double> > vSpace_;
   Teuchos::RCP<Thyra::VectorBase<double> > tVector_;
   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > tComm_;
};

template < >
class Response<panzer::Traits::Jacobian> : public ResponseBase {
public:
  Response(const std::string & responseName) : ResponseBase(responseName) 
  {
    tComm_ = Teuchos::DefaultComm<Thyra::Ordinal>::getComm();
  }

   //! What is the number of values you need locally
   virtual std::size_t localSizeRequired() const = 0;

   //! Is the vector distributed (or replicated). For derivative assembly this must be false!
   virtual bool vectorIsDistributed() const = 0;

    //! Does this response support derivative evaluation?
   bool supportsDerivative() const { return getDerivativeVectorSpace()!=Teuchos::null; }

   /** Get the vector for this response. This must be
     * constructed from the vector space returned by <code>getMap</code>.
     */
   Teuchos::RCP<Thyra::MultiVectorBase<double> > getDerivative() const
   { return derivative_; }

  //! Get the <code>Epetra_Map</code> for this response, map is constructed lazily.
   virtual Teuchos::RCP<Thyra::MultiVectorBase<double> > buildDerivative() const
   {
     TEUCHOS_ASSERT(!vectorIsDistributed());
     TEUCHOS_ASSERT(localSizeRequired()==1);
     TEUCHOS_ASSERT(supportsDerivative());
     return Thyra::createMember(*getDerivativeVectorSpace());
   }

   /** Set the vector (to be filled) for this response. This must be
     * constructed from the vector space returned by <code>getMap</code>.
     */
   virtual void setDerivative(const Teuchos::RCP<Thyra::MultiVectorBase<double> > & derivative)
   {
     TEUCHOS_ASSERT(!vectorIsDistributed());
     TEUCHOS_ASSERT(localSizeRequired()==1);
     TEUCHOS_ASSERT(supportsDerivative());
     derivative_ = derivative;
   }

protected:
   //! Get the teuchos comm object
   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > getComm() const { return tComm_; }

   //! Get the derivative vector space
   Teuchos::RCP<const Thyra::VectorSpaceBase<double> > getDerivativeVectorSpace() const
   { return derivVecSpace_; }

   //! Set the derivative vector space
   void setDerivativeVectorSpace(const Teuchos::RCP<const Thyra::VectorSpaceBase<double> > & vs)
   { derivVecSpace_ = vs; }

private:
   // hide these methods
   Response();
   Response(const Response<panzer::Traits::Jacobian> &);

   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > tComm_;
   Teuchos::RCP<const Thyra::VectorSpaceBase<double> > derivVecSpace_;

   Teuchos::RCP<Thyra::MultiVectorBase<double> > derivative_;
};

template < >
class Response<panzer::Traits::Tangent> : public ResponseBase {
public:
  Response(const std::string & responseName) : ResponseBase(responseName)
  {
    tComm_ = Teuchos::DefaultComm<Thyra::Ordinal>::getComm();
  }

   //! What is the number of values you need locally
   virtual std::size_t localSizeRequired() const = 0;

   //! Is the vector distributed (or replicated)
   virtual bool vectorIsDistributed() const = 0;

  //! Get the vector space for this response, vector space is constructed lazily.
   Teuchos::RCP<const Thyra::VectorSpaceBase<double> > getVectorSpace() const {
     // lazily build the space and return it
     if(vSpace_==Teuchos::null) {
       if(this->vectorIsDistributed())
         vSpace_ = Thyra::defaultSpmdVectorSpace<double>(tComm_,this->localSizeRequired(),-1);
       else
         vSpace_ = Thyra::locallyReplicatedDefaultSpmdVectorSpace<double>(tComm_,this->localSizeRequired());
     }
     return vSpace_;
   }

   /** Set the vector (to be filled) for this response. This must be
     * constructed from the vector space returned by <code>getVectorSpace</code>.
     */
   void setVector(const Teuchos::RCP<Thyra::MultiVectorBase<double> > & destVec) 
   {
     tVector_ = destVec;
   }

protected:
   //! Get the teuchos comm object
   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > getComm() const { return tComm_; }

   //! Access the thyra vector
   Thyra::ArrayRCP< Thyra::ArrayRCP<double> > getThyraMultiVector() const {
     const int num_col = tVector_->domain()->dim();
     Thyra::ArrayRCP< Thyra::ArrayRCP<double> > data(num_col);
     for (int i=0; i<num_col; ++i)
       Teuchos::rcp_dynamic_cast<Thyra::SpmdVectorBase<double> >(tVector_->col(i),true)->getNonconstLocalData(Teuchos::outArg(data[i]));
     return data;
  }

  //! Return the number of columns in the multivector
  int numDeriv() const {
      return tVector_->domain()->dim();
  }

private:
   // hide these methods
   Response();
   Response(const Response<panzer::Traits::Tangent> &);

   mutable Teuchos::RCP<const Thyra::VectorSpaceBase<double> > vSpace_;
   Teuchos::RCP<Thyra::MultiVectorBase<double> > tVector_;
   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > tComm_;
};

#ifdef Panzer_BUILD_HESSIAN_SUPPORT
template < >
class Response<panzer::Traits::Hessian> : public ResponseBase {
public:
   Response(const std::string & responseName) : ResponseBase(responseName) 
   {
    tComm_ = Teuchos::DefaultComm<Thyra::Ordinal>::getComm();
  }

   //! What is the number of values you need locally
   virtual std::size_t localSizeRequired() const = 0;

   //! Is the vector distributed (or replicated). For derivative assembly this must be false!
   virtual bool vectorIsDistributed() const = 0;

   //! Does this response support derivative evaluation?
   bool supportsDerivative() const { return getDerivativeVectorSpace()!=Teuchos::null; }

   //! Does this response support derivative evaluation?
   bool supportsDerivative() const { return getDerivativeVectorSpace()!=Teuchos::null; }

   /** Get the vector for this response. This must be
     * constructed from the vector space returned by <code>getMap</code>.
     */
   Teuchos::RCP<Thyra::MultiVectorBase<double> > getDerivative() const
   { return derivative_; }

   //! Get the <code>Epetra_Map</code> for this response, map is constructed lazily.
   virtual Teuchos::RCP<Thyra::MultiVectorBase<double> > buildDerivative() const
   {
     TEUCHOS_ASSERT(!vectorIsDistributed());
     TEUCHOS_ASSERT(localSizeRequired()==1);
     TEUCHOS_ASSERT(supportsDerivative());
     return Thyra::createMember(*getDerivativeVectorSpace());
   }

   /** Set the vector (to be filled) for this response. This must be
     * constructed from the vector space returned by <code>getMap</code>.
     */
   virtual void setDerivative(const Teuchos::RCP<Thyra::MultiVectorBase<double> > & derivative)
   {
     TEUCHOS_ASSERT(!vectorIsDistributed());
     TEUCHOS_ASSERT(localSizeRequired()==1);
     TEUCHOS_ASSERT(supportsDerivative());
     derivative_ = derivative;
   }

protected:
   //! Get the teuchos comm object
   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > getComm() const { return tComm_; }

   //! Get the derivative vector space
   Teuchos::RCP<const Thyra::VectorSpaceBase<double> > getDerivativeVectorSpace() const
   { return derivVecSpace_; }

   //! Set the derivative vector space
   void setDerivativeVectorSpace(const Teuchos::RCP<const Thyra::VectorSpaceBase<double> > & vs)
   { derivVecSpace_ = vs; }

private:
   // hide these methods
   Response();
   Response(const ResponseMESupportBase<panzer::Traits::Hessian> &);

   Teuchos::RCP<const Teuchos::Comm<Thyra::Ordinal> > tComm_;
   Teuchos::RCP<const Thyra::VectorSpaceBase<double> > derivVecSpace_;

   Teuchos::RCP<Thyra::MultiVectorBase<double> > derivative_;
};

#endif

}

#include "Response_impl.hpp"

#endif
